package com.example.studyllm.component.model;

import cn.hutool.core.util.StrUtil;
import com.example.studyllm.bo.ModelBO;
import com.example.studyllm.config.exception.BusinessException;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import lombok.extern.slf4j.Slf4j;

/**
 * @author wk
 */
@Slf4j
public class OpenAiChatModelBuilderHandler implements ChatModelBuilder{
    @Override
    public void checkParam(ModelBO modelBO) {
        if(StrUtil.isEmpty(modelBO.getApiKey())){
            throw new BusinessException("请上传该模型的appKey");
        }
    }

    @Override
    public ChatLanguageModel buildModel(ModelBO modelBO) {
        try {
            return OpenAiChatModel
                    .builder()
                    .apiKey(modelBO.getApiKey())
                    .baseUrl(modelBO.getBaseUrl())
                    .modelName(modelBO.getModel())
                    .maxTokens(modelBO.getResponseLength())
                    .temperature(modelBO.getTemperature())
                    .logRequests(true)
                    .logResponses(true)
                    .topP(modelBO.getTop())
                    .build();
        } catch (Exception e) {
            log.error("openai配置报错", e);
            return null;
        }
    }

    @Override
    public StreamingChatLanguageModel buildStreamModel(ModelBO modelBO) {
        try {
            return OpenAiStreamingChatModel
                    .builder()
                    .apiKey(modelBO.getApiKey())
                    .baseUrl(modelBO.getBaseUrl())
                    .modelName(modelBO.getModel())
                    .maxTokens(modelBO.getResponseLength())
                    .temperature(modelBO.getTemperature())
                    .logRequests(true)
                    .logResponses(true)
                    .topP(modelBO.getTop())
                    .build();
        } catch (Exception e) {
            log.error("openai配置报错", e);
            return null;
        }
    }

    @Override
    public EmbeddingModel buildEmbeddingModel(ModelBO modelBO) {
        try {
            OpenAiEmbeddingModel openAiEmbeddingModel = OpenAiEmbeddingModel
                    .builder()
                    .apiKey(modelBO.getApiKey())
                    .baseUrl(modelBO.getBaseUrl())
                    .modelName(modelBO.getModel())
                    .dimensions(modelBO.getDimension())
                    .logRequests(true)
                    .logResponses(true)
                    .dimensions(1024)
                    .build();
            return openAiEmbeddingModel;
        }catch (Exception e) {
            log.error(modelBO.getSupplier() + " Embedding 模型配置报错", e);
            return null;
        }
    }
}
